'''
Created on 02/06/2010

@author: Glazner
'''
from Queue import Queue, Empty
import unittest
from pymock import PyMockTestCase
from wp import Worker,WorkerPoolEvents,WorkerPool
from multiprocessing import Process

import time
import os
COMPUTER_NAME = os.getenv("COMPUTERNAME")

class MyWorker(Worker):
    def long_func(self,time2wait):
        time.sleep(time2wait)
        import os
        return os.getenv("COMPUTERNAME")

class MyPoolEvents(WorkerPoolEvents):
    def __init__(self):
        self.calls = Queue()
        
    def onHostConnected(self,host):
        self.calls.put(('onHostConnected',host))
    
    def onHostConnectionLost(self,host,exception):
        self.calls.put(('onHostConnectionLost',host,exception))
    
    def onHostIdle(self,host):
        self.calls.put(('onHostIdle',host))
    
    def onHostStartWorking(self,host,task):
        self.calls.put(('onHostStartWorking',host,task))
 
class TestWorkerPool(PyMockTestCase):
    
    def setUp(self):
        PyMockTestCase.setUp(self)
        self.processes = []
        p = Process(target=MyWorker().start,args=(50000,))
        p.start()
        self.processes.append(p)
        p = Process(target=MyWorker().start,args=(50001,))
        p.start()
        self.processes.append(p)
        self.poolEvents = MyPoolEvents()
        self.pool = WorkerPool(('localhost','localhost:50001'), self.poolEvents)
        time.sleep(1)
        
    def tearDown(self):
        PyMockTestCase.tearDown(self)
        for p in self.processes: p.terminate()
        self.pool.close()
        
    def testSubmit(self):
        l = []
        t = time.time()
        self.pool.submit('long_func',[5], lambda res: l.append(res))
        self.pool.submit('long_func',[4], lambda res: l.append(res))
        self.pool.wait()
        totalTime = time.time()-t
        assert 9>totalTime>5,totalTime
        assert l == [COMPUTER_NAME,COMPUTER_NAME]
        calls = []
        try:
            while True:
                calls.append(self.poolEvents.calls.get(False))
        except Empty:pass
        d = {'localhost':set(),'localhost:50001':set()}
        for data in calls:
            d[data[1]].add(data[0])
        self.assertEqual( d['localhost'],
            set(['onHostIdle', 'onHostConnected', 'onHostStartWorking']))
        self.assertEqual( d['localhost:50001'],
            set(['onHostIdle', 'onHostConnected', 'onHostStartWorking']))
    
    def testWait(self):
        for i in xrange(20):
            if len(self.pool.getActiveNodes()) == 2: break
            time.sleep(1) #lets servers power up
        t = time.time()
        self.pool.submit('long_func',[5])
        self.pool.submit('long_func',[4],group='Job1')
        self.pool.wait('Job1')
        totalTime = time.time()-t
        assert 5>totalTime>4,totalTime
        
    def testStop(self):
        t =time.time()
        self.pool.submit('long_func',[5])
        time.sleep(1)#let the first job pick up
        self.pool.submit('long_func',[5])
        self.pool.stop()
        self.pool.wait()
        total = time.time() - t
        assert 5 < total < 9 ,total
        
if __name__ == "__main__":
    unittest.main()